#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Feb 17 11:44:57 2024

@author: jcfq2
"""

import os

jwst_dir='/Users/jcfq2/data/observations/jwst'

os.chdir(jwst_dir)

import matplotlib.pyplot as plt
from astropy.visualization import make_lupton_rgb
from astropy.table import Table
import ch4_fiddlesticks as ch4
import pandas as pd
import h3ppy
import glob
from scipy.optimize import curve_fit
import spiceypy as spice
from astropy.io import fits
import numpy as np
from importlib import reload
import JWSTSolarSystemPointing as jssp
from scipy.interpolate import interp1d

reload(jssp)
#import jwst_uranus as jwstu

# In[0]

kernel_dir = '/Users/jcfq2/data/observations/jwst/kernels/'
jssp.load_kernels(kdir=kernel_dir)

# Set up a h3ppy object too, always useful
h3p_model = h3ppy.h3p()

# In[1]

fn = 'jw05308001001_03119_g395h-f290lp_s3d.fits'
file_dir = '/Users/jcfq2/data/observations/jwst/5308_dither_combined/'

file = file_dir + fn
# Create a JWSTSolarSystemPointing object that helps with lots of things JWST
geo = jssp.JWSTSolarSystemPointing(file)

# Caluclate the geometry for the full IFU cube. This is a three dimensional arrray.
cube = geo.full_fov()

# Get the wavelength scale for this observation - should be the same for all NIRSpec G395H/F290L observations.
wave = geo.get_wavelength()

# These are the available geometric parameters - if you need something that's not here, let me know!
# Only using Pandas to make this table pretty, generally not a fan
pd.DataFrame(geo.keys, columns=["Parameter"])



# %% make an intensity map using the PSG subtraction

import ch4_fiddlesticks as ch4

files = sorted(glob.glob(
    "/Users/jcfq2/data/observations/jwst/5308_dither_separated/*.fits"))
geo = jssp.JWSTSolarSystemPointing(files[0])
wave = geo.get_wavelength()
#  just to get geo for conversion

nam=np.load('saturn_spectra_map_2.npy')
# namlos=np.load('saturn_spectra_los_shell_2.npy')
nac=np.load('saturn_spectra_count_2.npy')

scale=2




# %%
# for ww in range(len(nam[0,0,:])):
    # nam[:,:,ww]=nam[:,:,ww]*namlos
# nam=nam*namlos
# namlos=0.0




nam[-90*scale:-45*scale,:]=nam[-90*scale:-45*scale,:]+nam[0:45*scale,:]
nac[-90*scale:-45*scale,:]=nac[-90*scale:-45*scale,:]+nac[0:45*scale,:]
nam[45*scale:90*scale,:]=nam[45*scale:90*scale,:]+nam[-45*scale:,:]
nac[45*scale:90*scale,:]=nac[45*scale:90*scale,:]+nac[-45*scale:,:]

# plt.imshow(nam[:,:,600]/nac[:,:,600])
# plt.show()



nam=nam[45*scale:-45*scale,:]
nac=nac[45*scale:-45*scale,:]

spec_cube=nam/nac
spec_cube=np.nan_to_num(spec_cube)

# %%

fig = plt.figure(figsize=(12,8),dpi=300)

numcount=plt.imshow(np.rot90(nac[:,220:,500]),cmap='tab20c')
cbar=fig.colorbar(numcount,location='bottom',aspect=4)

plt.show()



# %%


# nam=0.0
nac=0.0

# could be problematic - lots going on in the background 
file_331='/Users/jcfq2/data/observations/jwst/saturn/PSG_files/psg_rad_331.txt'
bck_331 = np.loadtxt(file_331)


# could be problematic - lots going on in the background 
file_345='/Users/jcfq2/data/observations/jwst/saturn/PSG_files/psg_rad_345.txt'
bck_345 = np.loadtxt(file_345)



# could be problematic - lots going on in the background 
file_353='/Users/jcfq2/data/observations/jwst/saturn/PSG_files/psg_rad_353.txt'
bck_353 = np.loadtxt(file_353)


file_360='/Users/jcfq2/data/observations/jwst/saturn/PSG_files/psg_rad_360_363.txt'
bck_360 = np.loadtxt(file_360)



file_366='/Users/jcfq2/data/observations/jwst/saturn/PSG_files/psg_rad_366_3675.txt'
bck_366 = np.loadtxt(file_366)


#

file_371_372='/Users/jcfq2/data/observations/jwst/saturn/PSG_files/psg_rad_371_372.txt'
bck_371_372 = np.loadtxt(file_371_372)


# horrid background
# file_388_389='/Users/jcfq2/data/observations/jwst/saturn/PSG_files/psg_rad_388_389.txt'
# bck_388_389 = np.loadtxt(file_388_389)

file_390_391='/Users/jcfq2/data/observations/jwst/saturn/PSG_files/psg_rad_390_391.txt'
bck_390_391 = np.loadtxt(file_390_391)


file_392_3935='/Users/jcfq2/data/observations/jwst/saturn/PSG_files/psg_rad_392_3935.txt'
bck_392_3935 = np.loadtxt(file_392_3935)

file_393_3958='/Users/jcfq2/data/observations/jwst/saturn/PSG_files/psg_rad_3930_3959.txt'
bck_393_3958 = np.loadtxt(file_393_3958)


file_396_398='/Users/jcfq2/data/observations/jwst/saturn/PSG_files/psg_rad_396_398.txt'
bck_396_398 = np.loadtxt(file_396_398)


file_398_399='/Users/jcfq2/data/observations/jwst/saturn/PSG_files/psg_rad_398_399.txt'
bck_398_399 = np.loadtxt(file_398_399)


file = files[-2]
# Create a JWSTSolarSystemPointing object that helps with lots of things JWST
geo = jssp.JWSTSolarSystemPointing(file)

# Caluclate the geometry for the full IFU cube. This is a three dimensional arrray.
cube = geo.full_fov()

# Get the wavelength scale for this observation - should be the same for all NIRSpec G395H/F290L observations.
wave = geo.get_wavelength()

# this didn't work with PSG, and is outside the range used by CH4
wavemin_331 = 3.31
wavemax_331 = 3.34
whw_331 = np.argwhere((wave > wavemin_331) & (wave < wavemax_331)).flatten()


wavemin_345 = 3.445
wavemax_345 = 3.465
whw_345 = np.argwhere((wave > wavemin_345) & (wave < wavemax_345)).flatten()


wavemin_353 = 3.5255
wavemax_353 = 3.565
whw_353 = np.argwhere((wave > wavemin_353) & (wave < wavemax_353)).flatten()


wavemin_360 = 3.61
wavemax_360 = 3.635
whw_360 = np.argwhere((wave > wavemin_360) & (wave < wavemax_360)).flatten()


wavemin_366 = 3.662
wavemax_366 = 3.675
whw_366 = np.argwhere((wave > wavemin_366) & (wave < wavemax_366)).flatten()


# correction for 3.9008

wavemin_371 = 3.709
wavemax_371 = 3.721
whw_371 = np.argwhere((wave > wavemin_371) & (wave < wavemax_371)).flatten()


# correction for 3.9008

wavemin_390 = 3.898
wavemax_390 = 3.911
whw_390 = np.argwhere((wave > wavemin_390) & (wave < wavemax_390)).flatten()


# correction for 3953

wavemin_392 = 3.923
wavemax_392 = 3.9313
whw_392 = np.argwhere((wave > wavemin_392) & (wave < wavemax_392)).flatten()


# correction for 3953

wavemin_3953 = 3.93
wavemax_3953 = 3.958
whw_3953 = np.argwhere((wave > wavemin_3953) & (wave < wavemax_3953)).flatten()

# correction for 3953
wavemin_396 = 3.961
wavemax_398 = 3.979
whw_3966 = np.argwhere((wave > wavemin_396) & (wave < wavemax_398)).flatten()

# correction for 3953
wavemin_398 = 3.98
wavemax_399 = 3.991
whw_398 = np.argwhere((wave > wavemin_398) & (wave < wavemax_399)).flatten()

# to fit temperature

# wavemin = 3.3
# wavemax = 4.3

# correction for 3953
wavemin_F323 = 3.148
wavemax_F323 = 3.35
whw_F323 = np.argwhere((wave > wavemin_F323) & (wave < wavemax_F323)).flatten()

F323_values=np.loadtxt('F323N_filter.txt',delimiter=' ')

f_F323 = interp1d(F323_values[:,0], F323_values[:,1], kind='linear')  # You can change 'linear' to other interpolation methods like 'cubic'


# whw = np.argwhere((wave > wavemin) & (wave < wavemax)).flatten()

h3p = h3ppy.h3p()
h3p.set(T=500, N=6e14, wave=wave, R=4700)
h3p_model=h3p.model()

im_h3p = np.nanmedian(geo.im[whw_3953, :, :], axis=0)
 
img = nam[:,:,1000]
nam=0.0


# Create an interpolation function
f_331 = interp1d(bck_331[:,0], bck_331[:,1], kind='linear')  # You can change 'linear' to other interpolation methods like 'cubic'
f_345 = interp1d(bck_345[:,0], bck_345[:,1], kind='linear')  # You can change 'linear' to other interpolation methods like 'cubic'
f_353 = interp1d(bck_353[:,0], bck_353[:,1], kind='linear')  # You can change 'linear' to other interpolation methods like 'cubic'
f_360 = interp1d(bck_360[:,0], bck_360[:,1], kind='linear')  # You can change 'linear' to other interpolation methods like 'cubic'
f_366 = interp1d(bck_366[:,0], bck_366[:,1], kind='linear')  # You can change 'linear' to other interpolation methods like 'cubic'
f_371 = interp1d(bck_371_372[:,0], bck_371_372[:,1], kind='linear')  # You can change 'linear' to other interpolation methods like 'cubic'
f_390 = interp1d(bck_390_391[:,0], bck_390_391[:,1], kind='linear')  # You can change 'linear' to other interpolation methods like 'cubic'
f_392 = interp1d(bck_392_3935[:,0], bck_392_3935[:,1], kind='linear')  # You can change 'linear' to other interpolation methods like 'cubic'
f_3953 = interp1d(bck_393_3958[:,0], bck_393_3958[:,1], kind='linear')  # You can change 'linear' to other interpolation methods like 'cubic'
f_3966 = interp1d(bck_396_398[:,0], bck_396_398[:,1], kind='linear')  # You can change 'linear' to other interpolation methods like 'cubic'
f_398 = interp1d(bck_398_399[:,0], bck_398_399[:,1], kind='linear')  # You can change 'linear' to other interpolation methods like 'cubic'

# Interpolate y values to the new x-range
bck_331_whw = f_331(wave[whw_331])
bck_345_whw = f_345(wave[whw_345])
bck_353_whw = f_353(wave[whw_353])
bck_360_whw = f_360(wave[whw_360])
bck_366_whw = f_366(wave[whw_366])
bck_371_372_whw = f_371(wave[whw_371])
bck_390_391_whw = f_390(wave[whw_390])
bck_392_3935_whw = f_392(wave[whw_392])
bck_393_3958_whw = f_3953(wave[whw_3953])
bck_396_398_whw = f_3966(wave[whw_3966])
bck_398_399_whw = f_398(wave[whw_398])
filter_F323_sub = f_F323(wave[whw_F323])

filter_F323 = np.zeros_like(wave)
filter_F323[whw_F323]=filter_F323_sub

# planetmask=np.zeros_like(img)
# temperature = np.zeros_like(img)
intensity = np.zeros_like(img)
intensity = intensity[:,120*scale:]
temperature = np.zeros_like(intensity)
pre_temperature = np.zeros_like(intensity)
temperature_error = np.zeros_like(intensity)
density = np.zeros_like(intensity)
pre_density = np.zeros_like(intensity)
density_error = np.zeros_like(intensity)
totalE = np.zeros_like(intensity)
pre_totalE = np.zeros_like(intensity)

ch4_fun = np.zeros_like(intensity)
ch4_hot = np.zeros_like(intensity)


# temperature2 = np.zeros_like(img)
# density = np.zeros_like(img)
# density2 = np.zeros_like(img)
# temperature_error = np.zeros_like(img)
# density_error = np.zeros_like(img)


# np.save('saturn_h3p_intensity.npy',intensity)

xs = 45*scale+1
ys = 360*scale+1

# planetmask[rayheight < 2000]=1
# planetmask[rayheight<-10000] = 0

# planetmask[:,0:2]=0
# planetmask[:,-1]=0


# Set a wavelength range where the non-LTE CH4 is
# whw = np.argwhere((wave > 3.9) & (wave < 4.6)).flatten()
# print(whw_3953)
# Fit the non-LTE CH4 background




        # for i in range(15) :             
        #     if (i == 0) : whl = np.argwhere((wave > 3.29) & (wave < 3.31)).flatten()
        #     if (i == 1) : whl = np.argwhere((wave > 3.37) & (wave < 3.4)).flatten()
        #     if (i == 2) : whl = np.argwhere((wave > 3.4) & (wave < 3.43)).flatten()
        #     if (i == 3) : whl = np.argwhere((wave > 3.445) & (wave < 3.465)).flatten()
        #     if (i == 4) : whl = np.argwhere((wave > 3.526) & (wave < 3.56)).flatten()
        #     if (i == 5) : whl = np.argwhere((wave > 3.60) & (wave < 3.63)).flatten()
        #     if (i == 6) : whl = np.argwhere((wave > 3.662) & (wave < 3.672)).flatten()
        #     if (i == 7) : whl = np.argwhere((wave > 3.709) & (wave < 3.72)).flatten()
        #   #  if (i == 8) : whl = np.argwhere((wave > 3.862) & (wave < 3.88)).flatten()
        #   #  if (i == 9) : whl = np.argwhere((wave > 3.8762) & (wave < 3.893)).flatten()
        #     if (i == 8) : whl = np.argwhere((wave > 3.9008) & (wave < 3.911)).flatten() 
        #     if (i == 9) : whl = np.argwhere((wave > 3.9235) & (wave < 3.9318)).flatten() 
        #     if (i == 10) : whl = np.argwhere((wave > 3.9495) & (wave < 3.957)).flatten() 
        #     if (i == 11) : whl = np.argwhere((wave > 3.9666) & (wave < 3.9775)).flatten() 
        #     if (i == 12) : whl = np.argwhere((wave > 3.98) & (wave < 3.99)).flatten() 
        #     if (i == 13) : whl = np.argwhere((wave > 4.18) & (wave < 4.24)).flatten() 
        #     if (i == 14) : whl = np.argwhere((wave > 4.34) & (wave < 4.4)).flatten() 




# %%




# temperature=np.load('saturn_saves/saturn_temperature_2z.npy')
# pre_temperature=np.load('saturn_saves/saturn_pretemperature_2z.npy')
# temperature_error=np.load('saturn_saves/saturn_temperature_error_2z.npy')
# density=np.load('saturn_saves/saturn_density_2z.npy')
# pre_density=np.load('saturn_saves/saturn_predensity_2z.npy')
# density_error=np.load('saturn_saves/saturn_density_error_2z.npy')
# totalE=np.load('saturn_saves/saturn_totalE_2z.npy')
# intensity=np.load('saturn_saves/saturn_intensity_2z.npy')
# ch4_fun=np.load('saturn_saves/saturn_ch4fun_2z.npy')
# ch4_hot=np.load('saturn_saves/saturn_ch4hot_2z.npy')




# cheats in the earlier fits
wavemin_331 = 3.27
wavemax_331 = 3.47
whw_331 = np.argwhere((wave > wavemin_331) & (wave < wavemax_331)).flatten()


# cheats in the earlier fits
wavemin_4 = 3.97
wavemax_4 = 3.99
whw_4 = np.argwhere((wave > wavemin_4) & (wave < wavemax_4)).flatten()


old_temp = ch4_hot
# #accounts for breaks and allowws reruns where data has big errors
# old_temp=np.load('saturn_saves/saturn_temperature_2.npy')+np.load('saturn_saves/saturn_temperature_2b.npy')
# old_temp_error=np.load('saturn_saves/saturn_temperature_error_2.npy')+np.load('saturn_saves/saturn_temperature_error_2b.npy')
# old_temp_ratio=old_temp_error/old_temp

# old_temp[old_temp_ratio >0.06]=0
# # old_temp[old_temp_ratio >0.0]=0

temp_ratio_factors=np.array([198.58633939662195,227.15566799054042,37.28510863758654])
rq_int_ratio_factors = np.array([-31.121753993057936,0.1082841194633575,-0.0001632594032538115,9.08369715525734e-08])




# degrees above 45: 30 = 75
startx=0*scale



for xxx in range(startx,xs):
    for yyy in range(ys):
        
# for xxx in range(2):
#     print(xxx)
#     for yyy in range(2):
        xx=xxx+135*scale#(left/right)
        yy=yyy#(up/down)
        xx_x = xx-135*scale

        # xx=xxx+160*scale#(left/right)
        # yy=yyy+20#(up/down)
        # if planetmask[yy,xx] > 0:
        if old_temp[yy,xx_x] == 0:
            spec = geo.convert(wave, spec_cube[yy, xx,:])  #auroral
            if np.sum(spec) == 0:
                print(xx/scale,yy/scale)
            else:
    
                # spec_old = geo.convert(wave, spec_cube[:, yy, xx])  #auroral
                # spec = geo.convert(wave, geo.im[:, 45, 34])  #auroral
                # spec = geo.convert(wave, geo.im[:, 40, 10])  #non-auroral
                
                # such a mess? doesn't match PSG at all - methane fluor?
                #3.66
                # bck_345_whw = f_345(wave[whw_345])
                # bck_345_whw=bck_345_whw/np.quantile(bck_345_whw[-10:],0.75)*(np.quantile(spec[whw_345][-10:],0.75))
                # # bck_345_whw=bck_345_whw/np.quantile(bck_345[:,1],0.85)*np.quantile(spec[whw_345],0.85)
                # spec[whw_345]=spec[whw_345]-bck_345_whw
        
                # PSG completely rubbish here...
                #3.66
                # bck_331_whw = f_331(wave[whw_331])
                # bck_331_whw=bck_331_whw/np.quantile(bck_331_whw[-10:],0.75)*(np.quantile(spec[whw_331][-10:],0.75))
                # bck_331_whw=bck_331_whw/np.quantile(bck_331[:,1],0.85)*np.quantile(spec[whw_331],0.85)
                spec[whw_331]=spec[whw_331]#-bck_331_whw
        
                
        
                #3.66
                bck_353_whw = f_353(wave[whw_353])
                bck_353_whw=bck_353_whw/np.quantile(bck_353_whw[-10:],0.75)*(np.quantile(spec[whw_353][-10:],0.75))
                # bck_353_whw=bck_353_whw/np.quantile(bck_353[:,1],0.85)*np.quantile(spec[whw_353],0.85)
                spec[whw_353]=spec[whw_353]-bck_353_whw
        
        
                #3.66
                bck_360_whw = f_360(wave[whw_360])
                bck_360_whw=bck_360_whw/np.quantile(bck_360_whw[-10:],0.75)*(np.quantile(spec[whw_360][-10:],0.75))
                # bck_360_whw=bck_360_whw/np.quantile(bck_360[:,1],0.85)*np.quantile(spec[whw_360],0.85)
                spec[whw_360]=spec[whw_360]-bck_360_whw
        
                #3.66
                bck_366_whw = f_366(wave[whw_366])
                bck_366_whw=bck_366_whw/np.quantile(bck_366_whw[-9:],0.75)*(np.quantile(spec[whw_366][-9:],0.75)-np.quantile(spec[whw_366][:5],0.5))
                # bck_366_whw=bck_366_whw/np.quantile(bck_366[:,1],0.85)*np.quantile(spec[whw_366],0.85)
                spec[whw_366]=spec[whw_366]-bck_366_whw
        
        
                #3.71
                bck_371_372_whw = f_371(wave[whw_371])
                bck_371_372_whw=bck_371_372_whw/np.quantile(bck_371_372_whw[-5:],0.75)*(np.quantile(spec[whw_371][-5:],0.75))
                # bck_371_372_whw=bck_371_372_whw/np.quantile(bck_371_372[:,1],0.5)*np.quantile(spec[whw_388],0.55)
                spec[whw_371]=spec[whw_371]-bck_371_372_whw
        
                # # lets make some background for h3ppy to fit no h3p to
                # whw_null=np.arange(500)
                # wave_null=wave[whw_null]
                # spec_null = np.zeros_like(wave_null)+np.random.normal(0, 1e-7, wave_null.shape) 
        
                # # lets make some background for h3ppy to fit no h3p to
                # whw_null2=np.arange(150)+1300
                # wave_null2=wave[whw_null2]
                # spec_null2 = np.zeros_like(wave_null2)+np.random.normal(0, 1e-7, wave_null2.shape) 
        
                
                #3.90
                bck_390_391_whw = f_390(wave[whw_390])
                bck_390_391_whw=bck_390_391_whw/np.quantile(bck_390_391_whw[0:6],0.75)*(np.quantile(spec[whw_390][0:6],0.75)-np.quantile(spec[whw_390][-4:],0.5))
                spec[whw_390]=spec[whw_390]-bck_390_391_whw
        
        
                #3.92
                bck_392_3935_whw = f_392(wave[whw_392])
                bck_392_3935_whw=bck_392_3935_whw/np.quantile(bck_392_3935[:,1],0.75)*np.quantile(spec[whw_392][0:6],0.75)
                spec[whw_392]=spec[whw_392]-bck_392_3935_whw
        
                
                #3.953
                bck_393_3958_whw = f_3953(wave[whw_3953])
                bck_393_3958_whw=bck_393_3958_whw/np.quantile(bck_393_3958[:,1],0.75)*np.quantile(spec[whw_3953],0.75)
                spec[whw_3953]=spec[whw_3953]-bck_393_3958_whw
        
                # not used - just imposible to pull out a signal from teh background
                # #3.966
                # bck_396_398_whw = f_3966(wave[whw_3966])
                # bck_396_398_whw=bck_396_398_whw/np.quantile(bck_396_398[:,1],0.75)*np.quantile(spec[whw_3966][0:10],0.75)
                # spec[whw_3966]=spec[whw_3966]-bck_396_398_whw
               
                    
                # not used -- to close to brightness edge to be trustable
                #3.98
                # bck_398_399_whw = f_398(wave[whw_398])
                # bck_398_399_whw=bck_398_399_whw/np.quantile(bck_398_399[:,1],0.75)*np.quantile(spec[whw_398],0.75)
                # spec[whw_398]=spec[whw_398]-bck_398_399_whw
        
        
                # spec_sub=np.hstack([spec_null,spec[whw_353],spec[whw_360],spec[whw_371],spec_null2,spec[whw_390],spec[whw_392],spec[whw_3953]])
                # whw_sub=np.hstack([whw_null,whw_353,whw_360,whw_371,whw_null2,whw_390,whw_392,whw_3953])

                # #  i useed this to finish the map
                spec_sub=np.hstack([spec[whw_331],spec[whw_353],spec[whw_360],spec[whw_371],spec[whw_390],spec[whw_392],spec[whw_3953]])
                whw_sub=np.hstack([whw_331,whw_353,whw_360,whw_371,whw_390,whw_392,whw_3953])

                # spec_sub=np.hstack([spec[whw_331],spec[whw_353],spec[whw_360],spec[whw_371],spec[whw_390],spec[whw_392],spec[whw_3953],spec[whw_4]])
                # whw_sub=np.hstack([whw_331,whw_353,whw_360,whw_371,whw_390,whw_392,whw_3953,whw_4])


                # I used this for the majority of the map
                # spec_sub=np.hstack([spec[whw_353],spec[whw_360],spec[whw_371],spec[whw_390],spec[whw_392],spec[whw_3953]])
                # whw_sub=np.hstack([whw_353,whw_360,whw_371,whw_390,whw_392,whw_3953])

                # spec_sub=np.hstack([spec[whw_353],spec[whw_360],spec[whw_371],spec[whw_390]])
                # whw_sub=np.hstack([whw_353,whw_360,whw_371,whw_390])

                wave_sub=wave[whw_sub]
                
                reload(ch4)
        
                ch4fit = ch4.fit_non_LTE_CH4_JWST(wave)
                fit = ch4fit.fit(wave, spec)
                # fit[whw_null]=0.
                # fit[whw_null2]=0.
        
                
                fit_sub = fit[whw_sub]
                
                fit_subA = fit[whw_353]
                fit_subB = fit[whw_3953]
                
                
                subwave = wave_sub[np.isfinite(fit_sub)]
                subspec = (spec_sub - fit_sub)[np.isfinite(fit_sub)]
                
                subback = fit_sub[np.isfinite(fit_sub)]
    
                subspecA = (spec[whw_353]-fit_subA)[np.isfinite(fit_subA)]
                subspecB = (spec[whw_3953]-fit_subB)[np.isfinite(fit_subB)]
    
                rq_ratio = np.sum(subspecA)/np.sum(subspecB)
                rq_temp=temp_ratio_factors[0]+temp_ratio_factors[1]*rq_ratio+temp_ratio_factors[2]*rq_ratio**2
    
                if rq_temp < 350: rq_temp=350
                if rq_temp > 550: rq_temp=550
    
                rq_int = np.sum(subspecA)+np.sum(subspecB)
                rq_scaling=np.exp(rq_int_ratio_factors[0]+rq_int_ratio_factors[1]*rq_temp+rq_int_ratio_factors[2]*rq_temp**2+rq_int_ratio_factors[3]*rq_temp**3)

                rq_density=rq_int/rq_scaling*1.5e16

                

                print(rq_ratio,rq_temp,rq_int,np.sum(spec[whw_353])/np.sum(spec[whw_3953]),rq_scaling,rq_density/1e15)
                # rq_temp=450
                # rq_density=rq_density*30
                h3p = h3ppy.h3p()
                # Fit the residual H3+ spectrum
                h3p.set(wave=subwave, data=subspec)
                h3p.set(T=rq_temp, N=rq_density)
                prefit_model = h3p.model()

                pre_temperature[yy,xx_x]=rq_temp
                pre_density[yy,xx_x]=rq_density
                pre_totalE[yy,xx_x]=h3p.total_emission()

                ratio_mask = np.zeros_like(prefit_model)+1.
                
                # if (np.sum(subspec[prefit_model > 0.1])/np.sum(ratio_mask[prefit_model > 0.1])) / (np.sum(subspec[prefit_model < 0.1])/np.sum(ratio_mask[prefit_model < 0.1])) > 10:
                intensity[yy,xx_x]=np.sum(subspec)
                ch4_fun[yy,xx_x]=np.nansum(ch4fit.ch4_fun[whw_sub]*ch4fit.fun_scaling)
                ch4_hot[yy,xx_x]=np.nansum(ch4fit.ch4_hot[whw_sub]*ch4fit.hot_scaling)
    
                            
                h3p = h3ppy.h3p()
                # Fit the residual H3+ spectrum
                h3p.set(wave=subwave, data=subspec)
                h3p.set(T=rq_temp, N=rq_density)
                # h3p.guess_density()
                h3p_fit = h3p.fit()
                vars, errs = h3p.get_results(verbose=False)
    
            # Fit the residual H3+ spectrum
    
                if vars == False: 
                    print(xx_x,xx/scale,yy/scale,' -Fit failed- ',round(pre_temperature[yy,xx_x], 3),', ',round(pre_density[yy,xx_x]/1e15, 3),', ',round(pre_totalE[yy,xx_x]/1e-6, 3) )
                else:
                    temperature[yy,xx_x]=vars['temperature']
                    density[yy,xx_x]=vars['density']
                    temperature_error[yy,xx_x]=errs['temperature']
                    density_error[yy,xx_x]=errs['density']
                    totalE[yy,xx_x]=h3p.total_emission()
                    
                    print( xx/scale,', ',yy/scale,', ',round(temperature[yy,xx_x], 3),'+/-',round(temperature_error[yy,xx_x], 2),', ',round(density[yy,xx_x]/1e15, 3),'+-',round(density_error[yy,xx_x]/1e15, 3),' ',round(totalE[yy,xx_x]/1e-6, 3) )
    
    
                fig, (ax1, ax2,ax3,ax4,ax5,ax6,ax7,ax8) = plt.subplots(1, 8, figsize=(12, 4), gridspec_kw={'width_ratios': [6,1,1, 1,1,1,1,1]})
                
                # ax1.plot(subwave,subspec,label='data')
                ax1.plot(subspec,label='data-model')
                ax1.plot(spec_sub,label='raw spec')
                # ax1.plot(wave[whw],spec[whw],label='raw spec')
                # ax1.plot(spec_sub[np.isfinite(fit_sub)]-fit_sub[np.isfinite(fit_sub)],label='data-fit')
                
                # ax1.plot(wave[whw],bck_353_whw,label='PSG scaled')
                # ax1.plot(wave[whw],spec[whw]+bck_353_whw,label='spec without PSG')
                # ax1.plot(wave[whw],fit[whw],label='back fit')
        
                # # ax1.plot(wave[whw],spec[whw]-fit[whw],label='data-fit')
                
                ax1.plot(ch4fit.ch4_fun[whw_sub]*ch4fit.fun_scaling,label='methane_fund',linestyle='dotted')
                ax1.plot(ch4fit.ch4_hot[whw_sub]*ch4fit.hot_scaling,label='methane_hot',linestyle='dotted')
                # ax1.plot(ch4fit.wave[whw_sub],(ch4fit.ch4_fun[whw_sub]+ch4fit.ch4_hot[whw_sub])*np.max(spec[whw_sub]),label='methane')
                # ax1.plot(h3p_model[whw_sub]/np.nanmax(h3p_model[whw_sub])*np.max(spec_sub),label='H3+',linestyle='dashed')
                # ax1.plot(mm,label='H3+ model'+str(temp_array[tt])+str(cd_array[dd]/1e15),linestyle='dashed')
                # ax1.plot(h3p_model[whw_sub][np.isfinite(fit_sub)],label='H3+ model',linestyle='dashed')
                ax1.plot(h3p_fit,label='H3+ fit',linestyle='dashed')
                ax1.plot(prefit_model,label='H3+ premodel',linestyle='dotted')
                # ax1.plot(wave,fit,label='Bkg fit')
                ax1.legend()
                ax1.set_xlim([0,subspec.size])
                # 
                testimage = np.nanmedian(spec_cube[:,:,whw_sub],axis=2)
                testimage[yy, xx] = np.nanmax(testimage)
                  
                ax2.imshow(ch4_fun+ch4_hot,vmin= 0)
                ax3.imshow(intensity,vmin= 0)#np.min(intensity[np.nonzero(intensity)]) )
                ax4.imshow(temperature,vmin= 0)#np.min(intensity[np.nonzero(intensity)]) )
                ax5.imshow(density**0.5,vmin= 0)#np.min(intensity[np.nonzero(intensity)]) )
                ax6.imshow(totalE**0.5,vmin= 0)#np.min(intensity[np.nonzero(intensity)]) )
                ax7.imshow(np.nan_to_num(temperature_error/temperature),vmin= 0,vmax=0.05)#np.min(intensity[np.nonzero(intensity)]) )
                ax8.imshow(np.nan_to_num(density_error/density),vmin= 0,vmax=0.2)#np.min(intensity[np.nonzero(intensity)]) )
        
                
                # plt.imshow(temperature)
                
                plt.show()
                
        else:
            print(xx/scale,yy/scale,'pre_fit')


np.save('saturn_saves/saturn_temperature_2z.npy',temperature)
np.save('saturn_saves/saturn_pretemperature_2z.npy',pre_temperature)
np.save('saturn_saves/saturn_temperature_error_2z.npy',temperature_error)
np.save('saturn_saves/saturn_density_2z.npy',density)
np.save('saturn_saves/saturn_predensity_2z.npy',pre_density)
np.save('saturn_saves/saturn_density_error_2z.npy',density_error)
np.save('saturn_saves/saturn_totalE_2z.npy',totalE)
np.save('saturn_saves/saturn_intensity_2z.npy',intensity)
np.save('saturn_saves/saturn_ch4fun_2z.npy',ch4_fun)
np.save('saturn_saves/saturn_ch4hot_2z.npy',ch4_hot)



#
## %%  split to save loading times during debug
#
#temperature=np.load('saturn_saves/saturn_temperature_2.npy')
#temperature_error=np.load('saturn_saves/saturn_temperature_error_2.npy')
#density=np.load('saturn_saves/saturn_density_2.npy')
#density_error=np.load('saturn_saves/saturn_density_error_2.npy')
#totalE=np.load('saturn_saves/saturn_totalE_2.npy')
#intensity=np.load('saturn_saves/saturn_intensity_2.npy')
#
## %%  split to save loading times during debug
#
#te=np.load('saturn_saves/saturn_temperature_2.npy')
#te_er=np.load('saturn_saves/saturn_temperature_error_2.npy')
#de=np.load('saturn_saves/saturn_density_2.npy')
#de_er=np.load('saturn_saves/saturn_density_error_2.npy')
#toE=np.load('saturn_saves/saturn_totalE_2.npy')
#inT=np.load('saturn_saves/saturn_intensity_2.npy')
#
#
#te2=np.load('saturn_saves/saturn_temperature_2b.npy')
#te_er2=np.load('saturn_saves/saturn_temperature_error_2b.npy')
#de2=np.load('saturn_saves/saturn_density_2b.npy')
#de_er2=np.load('saturn_saves/saturn_density_error_2b.npy')
#toE2=np.load('saturn_saves/saturn_totalE_2b.npy')
#inT2=np.load('saturn_saves/saturn_intensity_2b.npy')
#
#
#te3=np.load('saturn_saves/saturn_temperature_2c.npy')
#te_er3=np.load('saturn_saves/saturn_temperature_error_2c.npy')
#de3=np.load('saturn_saves/saturn_density_2c.npy')
#de_er3=np.load('saturn_saves/saturn_density_error_2c.npy')
#toE3=np.load('saturn_saves/saturn_totalE_2c.npy')
#inT3=np.load('saturn_saves/saturn_intensity_2c.npy')
#
#
#TE=te+te2
#TE[te3 > 0]=te3[te3 > 0]
#
#TE_ER = te_er+te_er2
#TE_ER[te3 > 0]=te_er3[te3 > 0]
#
#DE=de+de2
#DE[te3 > 0]=de3[te3 > 0]
#
#DE_ER = de_er+de_er2
#DE_ER[te3 > 0]=de_er3[te3 > 0]
#
#TOE = toE + toE2
#TOE[te3 > 0]=toE3[te3 > 0]
#
#INT = inT + inT2
#INT[te3 >0]=inT3[te3 > 0]
#
## %%  DO NOT USE!  copy code to end of file when complete - this is to guess temperatures and densities
#
#
#startx=58
#
#
#for xxx in range(1):
#    for yyy in range(1):
#        
## for xxx in range(2):
##     print(xxx)
##     for yyy in range(2):
#        xx=xxx+135*scale#(left/right)
#        yy=yyy#(up/down)
#        # xx=xxx+160*scale#(left/right)
#        # yy=yyy+20#(up/down)
#        # if planetmask[yy,xx] > 0:
#        spec = geo.convert(wave, spec_cube[yy, xx,:])  #auroral
#        if np.sum(spec) == 0:
#            print(xx/scale,yy/scale)
#        else:
#
#            # spec_old = geo.convert(wave, spec_cube[:, yy, xx])  #auroral
#            # spec = geo.convert(wave, geo.im[:, 45, 34])  #auroral
#            # spec = geo.convert(wave, geo.im[:, 40, 10])  #non-auroral
#            
#            # such a mess? doesn't match PSG at all - methane fluor?
#            #3.66
#            # bck_345_whw = f_345(wave[whw_345])
#            # bck_345_whw=bck_345_whw/np.quantile(bck_345_whw[-10:],0.75)*(np.quantile(spec[whw_345][-10:],0.75))
#            # # bck_345_whw=bck_345_whw/np.quantile(bck_345[:,1],0.85)*np.quantile(spec[whw_345],0.85)
#            # spec[whw_345]=spec[whw_345]-bck_345_whw
#    
#            # PSG completely rubbish here...
#            # #3.66
#            # bck_331_whw = f_331(wave[whw_331])
#            # bck_331_whw=bck_331_whw/np.quantile(bck_331_whw[-10:],0.75)*(np.quantile(spec[whw_331][-10:],0.75))
#            # # bck_331_whw=bck_331_whw/np.quantile(bck_331[:,1],0.85)*np.quantile(spec[whw_331],0.85)
#            # spec[whw_331]=spec[whw_331]-bck_331_whw
#    
#            
#    
#            #3.66
#            bck_353_whw = f_353(wave[whw_353])
#            bck_353_whw=bck_353_whw/np.quantile(bck_353_whw[-10:],0.75)*(np.quantile(spec[whw_353][-10:],0.75))
#            # bck_353_whw=bck_353_whw/np.quantile(bck_353[:,1],0.85)*np.quantile(spec[whw_353],0.85)
#            spec[whw_353]=spec[whw_353]-bck_353_whw
#    
#    
#            #3.66
#            bck_360_whw = f_360(wave[whw_360])
#            bck_360_whw=bck_360_whw/np.quantile(bck_360_whw[-10:],0.75)*(np.quantile(spec[whw_360][-10:],0.75))
#            # bck_360_whw=bck_360_whw/np.quantile(bck_360[:,1],0.85)*np.quantile(spec[whw_360],0.85)
#            spec[whw_360]=spec[whw_360]-bck_360_whw
#    
#            #3.66
#            bck_366_whw = f_366(wave[whw_366])
#            bck_366_whw=bck_366_whw/np.quantile(bck_366_whw[-9:],0.75)*(np.quantile(spec[whw_366][-9:],0.75)-np.quantile(spec[whw_366][:5],0.5))
#            # bck_366_whw=bck_366_whw/np.quantile(bck_366[:,1],0.85)*np.quantile(spec[whw_366],0.85)
#            spec[whw_366]=spec[whw_366]-bck_366_whw
#    
#    
#            #3.71
#            bck_371_372_whw = f_371(wave[whw_371])
#            bck_371_372_whw=bck_371_372_whw/np.quantile(bck_371_372_whw[-5:],0.75)*(np.quantile(spec[whw_371][-5:],0.75))
#            # bck_371_372_whw=bck_371_372_whw/np.quantile(bck_371_372[:,1],0.5)*np.quantile(spec[whw_388],0.55)
#            spec[whw_371]=spec[whw_371]-bck_371_372_whw
#    
#            # lets make some background for h3ppy to fit no h3p to
#            whw_null=np.arange(500)
#            wave_null=wave[whw_null]
#            spec_null = np.zeros_like(wave_null)+np.random.normal(0, 1e-7, wave_null.shape) 
#    
#            # lets make some background for h3ppy to fit no h3p to
#            whw_null2=np.arange(150)+1300
#            wave_null2=wave[whw_null2]
#            spec_null2 = np.zeros_like(wave_null2)+np.random.normal(0, 1e-7, wave_null2.shape) 
#    
#            
#            #3.90
#            bck_390_391_whw = f_390(wave[whw_390])
#            bck_390_391_whw=bck_390_391_whw/np.quantile(bck_390_391_whw[0:6],0.75)*(np.quantile(spec[whw_390][0:6],0.75)-np.quantile(spec[whw_390][-4:],0.5))
#            spec[whw_390]=spec[whw_390]-bck_390_391_whw
#    
#    
#            #3.92
#            bck_392_3935_whw = f_392(wave[whw_392])
#            bck_392_3935_whw=bck_392_3935_whw/np.quantile(bck_392_3935[:,1],0.75)*np.quantile(spec[whw_392][0:6],0.75)
#            spec[whw_392]=spec[whw_392]-bck_392_3935_whw
#    
#            
#            #3.953
#            bck_393_3958_whw = f_3953(wave[whw_3953])
#            bck_393_3958_whw=bck_393_3958_whw/np.quantile(bck_393_3958[:,1],0.75)*np.quantile(spec[whw_3953],0.75)
#            spec[whw_3953]=spec[whw_3953]-bck_393_3958_whw
#    
#            # not used - just imposible to pull out a signal from teh background
#            # #3.966
#            # bck_396_398_whw = f_3966(wave[whw_3966])
#            # bck_396_398_whw=bck_396_398_whw/np.quantile(bck_396_398[:,1],0.75)*np.quantile(spec[whw_3966][0:10],0.75)
#            # spec[whw_3966]=spec[whw_3966]-bck_396_398_whw
#           
#                
#            # not used -- to close to brightness edge to be trustable
#            #3.98
#            # bck_398_399_whw = f_398(wave[whw_398])
#            # bck_398_399_whw=bck_398_399_whw/np.quantile(bck_398_399[:,1],0.75)*np.quantile(spec[whw_398],0.75)
#            # spec[whw_398]=spec[whw_398]-bck_398_399_whw
#    
#    
#            # spec_sub=np.hstack([spec_null,spec[whw_353],spec[whw_360],spec[whw_371],spec_null2,spec[whw_390],spec[whw_392],spec[whw_3953]])
#            # whw_sub=np.hstack([whw_null,whw_353,whw_360,whw_371,whw_null2,whw_390,whw_392,whw_3953])
#            spec_sub=np.hstack([spec[whw_353],spec[whw_360],spec[whw_371],spec[whw_390],spec[whw_392],spec[whw_3953]])
#            whw_sub=np.hstack([whw_353,whw_360,whw_371,whw_390,whw_392,whw_3953])
#            wave_sub=wave[whw_sub]
#            
#            reload(ch4)
#    
#            ch4fit = ch4.fit_non_LTE_CH4_JWST(wave)
#            fit = ch4fit.fit(wave, spec)
#            # fit[whw_null]=0.
#            # fit[whw_null2]=0.
#    
#            
#            fit_sub = fit[whw_sub]
#            
#          
#            
#            subwave = wave_sub[np.isfinite(fit_sub)]
#            subspec = (spec_sub - fit_sub)[np.isfinite(fit_sub)]
#            
#            subback = fit_sub[np.isfinite(fit_sub)]
#
#            wave_A = wave[whw_353]
#            wave_B = wave[whw_3953]
#
#            temp_ratio = np.zeros([200])
#            den_scale = np.zeros([200])
#
#            for Tt in range(200):
#                h3p = h3ppy.h3p()
#                # Fit the residual H3+ spectrum
#                h3p.set(wave=wave_A, data=subspec)
#                h3p.set(T=350+Tt, N=1e16)
#                prefit_waveA = h3p.model()
#                # plt.plot(prefit_waveA)
#                # plt.show()
#                h3p.set(wave=wave_B, data=subspec)
#                prefit_waveB = h3p.model()
#                # plt.plot(prefit_waveB)
#                # plt.show()
#                
#                print(np.sum(prefit_waveA),np.sum(prefit_waveB))
#            
#
#                temp_ratio[Tt]=np.sum(prefit_waveA)/np.sum(prefit_waveB)
#                den_scale[Tt]=np.sum(prefit_waveA)+np.sum(prefit_waveB)
#            
#            plt.plot(np.arange(350,550),temp_ratio)
#            plt.show()
#            plt.plot(np.arange(350,550),den_scale)
#            plt.show()
#
## %% fit the variance curves
#
#def poly_curve(x, A, B, C):
#    y = A+B*x+C*x**2
#    return y
#
#def poly_curve3(x, A, B, C, D):
#    y = A+B*x+C*x**2+D*x**3
#    return y
#
#xdata=np.arange(350,550)
#yydata=den_scale
#ydata=np.log(den_scale)
#parameters, covariance = curve_fit(poly_curve3, xdata, ydata)
#
#fit_A = parameters[0]
#fit_B = parameters[1]
#fit_C = parameters[2]
#fit_D = parameters[3]
#print(fit_A)
#print(fit_B)
#print(fit_C)
#print(fit_D)
#
#fit_y = poly_curve3(xdata, fit_A, fit_B,fit_C,fit_D)
#plt.plot(xdata, yydata, 'o', label='data')
#plt.plot(xdata, np.exp(fit_y), '-', label='fit')
## plt.plot(np.exp(xdata), xxdata, '-', label='fit')
#plt.legend()
#
#temp_factors=np.array([-0.8405218573956371,0.004629978824503198,-1.414140496241531e-06])
#temp_ratio_factors=np.array([198.58633939662195,227.15566799054042,37.28510863758654])
#
#rq_ratio=0.7
#
## print(temp_ratio_factors[0]+temp_ratio_factors[1]*rq_ratio+temp_ratio_factors[2]*rq_ratio**2)
#
#rq_temp = 480
#
#rq_int_ratio_factors = np.array([-31.121753993057936,0.1082841194633575,-0.0001632594032538115,9.08369715525734e-08])
#
#print(np.exp(rq_int_ratio_factors[0]+rq_int_ratio_factors[1]*rq_temp+rq_int_ratio_factors[2]*rq_temp**2+rq_int_ratio_factors[3]*rq_temp**3))
#
#
#
#